#include "AccelerationStructureGeometry.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../../Templates/NewObject.h"
#include "../../Game/Process/RenderProcess.h"

AccelerationStructureGeometry::Builder::Builder() noexcept = default;
AccelerationStructureGeometry::Builder::~Builder() noexcept = default;
AccelerationStructureGeometry::Builder::Builder(AccelerationStructureGeometry::Builder const& rhs) noexcept = default;
AccelerationStructureGeometry::Builder::Builder(AccelerationStructureGeometry::Builder && rhs) noexcept = default;
AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::operator=(AccelerationStructureGeometry::Builder const& rhs) noexcept = default;
AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::operator=(AccelerationStructureGeometry::Builder && rhs) noexcept = default;

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetPrimitiveCount(uint32 primitiveCount) noexcept
{
	mImpl->rangeInfo.primitiveCount = primitiveCount;
	return *this;
}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetPrimitiveOffset(uint32 primitiveOffset) noexcept
{
	mImpl->rangeInfo.primitiveOffset = primitiveOffset;
	return *this;
}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetFirstVertex(uint32 firstVertex) noexcept
{
	mImpl->rangeInfo.firstVertex = firstVertex;
	return *this;
}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetTransformOffset(uint32 transformOffset) noexcept
{
	mImpl->rangeInfo.transformOffset = transformOffset;
	return *this;
}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetGeometryType(VkGeometryTypeKHR geometryType) noexcept
{
	mImpl->geometryType = geometryType;
	return *this;
}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetFlags(VkGeometryFlagsKHR flags) noexcept
{
	mImpl->flags = flags;
	return *this;
}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetVertexFormat(VkFormat vertexFormat) noexcept
{
	mImpl->triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
	mImpl->triangles.vertexFormat = vertexFormat;
	return *this;
}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetVertexStride(VkDeviceSize vertexStride) noexcept
{
	mImpl->triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
	mImpl->triangles.vertexStride = vertexStride;
	return *this;
}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetMaxVertex(uint32 maxVertex) noexcept
{
	mImpl->triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
	mImpl->triangles.maxVertex = maxVertex;
	return *this;
}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetVertexData(RefCountPtr<RHIBuffer> vertexData) noexcept
{
	mImpl->triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
	mImpl->vertexData = vertexData;
	mImpl->triangles.vertexData.deviceAddress = vertexData->GetDeviceAddress();
	return *this;
}

//AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetVertexData(void* hostAddress) noexcept
//{
//	mImpl->triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
//	mImpl->triangles.vertexData.hostAddress = hostAddress;
//	return *this;
//}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetIndexType(VkIndexType indexType) noexcept
{
	mImpl->triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
	mImpl->triangles.indexType = indexType;
	return *this;
}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetIndexData(RefCountPtr<RHIBuffer> indexData) noexcept
{
	mImpl->triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
	mImpl->indexData = indexData;
	mImpl->triangles.indexData.deviceAddress = indexData->GetDeviceAddress();
	return *this;
}

//AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetIndexData(void* hostAddress) noexcept
//{
//	mImpl->triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
//	mImpl->triangles.indexData.hostAddress = hostAddress;
//	return *this;
//}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetTransformData(RefCountPtr<RHIBuffer> transformData) noexcept
{
	mImpl->triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
	mImpl->transformData = transformData;
	mImpl->triangles.transformData.deviceAddress = transformData->GetDeviceAddress();
	return *this;
}

//AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetTransformData(void* hostAddress) noexcept
//{
//	mImpl->triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
//	mImpl->triangles.transformData.hostAddress = hostAddress;
//	return *this;
//}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetAabbStride(VkDeviceSize aabbStride) noexcept
{
	mImpl->aabbs.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_AABBS_DATA_KHR;
	mImpl->aabbs.stride = aabbStride;
	return *this;
}

AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetAabbData(RefCountPtr<RHIBuffer> aabbData) noexcept
{
	mImpl->aabbs.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_AABBS_DATA_KHR;
	mImpl->aabbData = aabbData;
	mImpl->aabbs.data.deviceAddress = aabbData->GetDeviceAddress();
	return *this;
}

//AccelerationStructureGeometry::Builder& AccelerationStructureGeometry::Builder::SetAabbData(void* hostAddress) noexcept
//{
//	mImpl->aabbs.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_AABBS_DATA_KHR;
//	mImpl->aabbs.data.hostAddress = hostAddress;
//	return *this;
//}

RefCountPtr<AccelerationStructureGeometry> AccelerationStructureGeometry::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<AccelerationStructureGeometry>(NewObject<AccelerationStructureGeometry>(*mImpl));
}

RefCountPtr<AccelerationStructureGeometry> AccelerationStructureGeometry::Builder::Build(RenderProcess* process, const std::string& name)
{
	assert(IsInRenderThread());
	AccelerationStructureGeometry* accelerationStructureGeometry = NewObject<AccelerationStructureGeometry>(*mImpl);
	process->AddAccelerationStructureGeometry(name, accelerationStructureGeometry);
	return RefCountPtr<AccelerationStructureGeometry>(accelerationStructureGeometry);
}

void AccelerationStructureGeometry::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemoveAccelerationStructureGeometry(mName, this);
		}
		DeleteObject<AccelerationStructureGeometry>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemoveAccelerationStructureGeometry(mName, this);
			}
			DeleteObject<AccelerationStructureGeometry>(this);
			});
	}
}

AccelerationStructureGeometry::AccelerationStructureGeometry(BuilderDetails & builderDetails)
{
	mPrimitiveCount = 0;

	mRangeInfo = builderDetails.rangeInfo;

	mAccelerationStructureGeometry = {};
	mAccelerationStructureGeometry.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR;
	mAccelerationStructureGeometry.geometryType = builderDetails.geometryType;
	mAccelerationStructureGeometry.flags = builderDetails.flags;

	if (builderDetails.triangles.sType != VK_STRUCTURE_TYPE_MAX_ENUM)
	{
		assert(builderDetails.vertexData.IsValid());
		mVertexData = builderDetails.vertexData;
		assert(builderDetails.indexData.IsValid());
		mIndexData = builderDetails.indexData;
		// assert(builderDetails.transformData.IsValid());
		mTransformData = builderDetails.transformData;

		mAccelerationStructureGeometry.geometry.triangles = builderDetails.triangles;

		switch (mAccelerationStructureGeometry.geometry.triangles.indexType)
		{
		case VK_INDEX_TYPE_UINT16:
			mPrimitiveCount = static_cast<uint32>(mIndexData->GetSize() / 2);
			break;
		case VK_INDEX_TYPE_UINT32:
			mPrimitiveCount = static_cast<uint32>(mIndexData->GetSize() / 4);
			break;
		case VK_INDEX_TYPE_UINT8_EXT:
			mPrimitiveCount = static_cast<uint32>(mIndexData->GetSize());
			break;
		default:
			assert(0);
			break;
		}
		// PrimitiveCount 为 三角形数量
		mPrimitiveCount /= 3;
	}

	if (builderDetails.aabbs.sType != VK_STRUCTURE_TYPE_MAX_ENUM)
	{
		mAccelerationStructureGeometry.geometry.aabbs = builderDetails.aabbs;

		assert(builderDetails.aabbData.IsValid());
		mAabbData = builderDetails.aabbData;
		// PrimitiveCount 为 aabb盒数量, stride 为 sizeof(aabb)
		mPrimitiveCount = static_cast<uint32>(mAabbData->GetSize() / mAccelerationStructureGeometry.geometry.aabbs.stride);
	}

	// 检查并且重设rangeinfo
	if (mRangeInfo.primitiveCount != 0)
	{
		assert(mRangeInfo.primitiveCount == mPrimitiveCount);
	}
	if (mRangeInfo.primitiveCount == 0 && mPrimitiveCount != 0)
	{
		mRangeInfo.primitiveCount = mPrimitiveCount;
	}
}

VkAccelerationStructureGeometryKHR AccelerationStructureGeometry::GetAccelerationStructureGeometry() const
{
	return mAccelerationStructureGeometry;
}

uint32 AccelerationStructureGeometry::GetPrimitiveCount() const
{
	return mPrimitiveCount;
}

VkAccelerationStructureBuildRangeInfoKHR* AccelerationStructureGeometry::GetRangeInfo()
{
	return &mRangeInfo;
}